# file:main.py
# 实现相关运算(这里主要是损失函数、优化器等):
import torch
# 导入神经网络模型:
from CNN import ConvNet
# 导入数据:
from Dataloader import get_train_data, test_val_data
# 导入训练、测试方法:
from train_test import train_model, test_model
"""超参数"""
batch_size = 64
epochs = 10
learning_rate = 0.001
momentum = 0.9

if torch.cuda.is_available():
    device = 'cuda'
else:
    device = 'cpu'
print(device)

# 实例化神经网络模型:
# cnn_model = ConvNet().cuda()
cnn_model = ConvNet().to('cpu')
# 定义损失函数(交叉熵):
loss_func = torch.nn.CrossEntropyLoss()
# 定义优化器(随机梯度下降):
sgd_opt = torch.optim.SGD(cnn_model.parameters(),
						  # 学习率:
						  lr=learning_rate,
						  # 动量:
						  momentum=momentum)

# 读取数据:
_, train_loader = get_train_data()
_, test_loader, val_loader = test_val_data()

# 开始训练:
train_model(cnn_model, train_loader, val_loader, batch_size, epochs, loss_func, sgd_opt)
test_model(cnn_model, test_loader)
# 保存模型，以便之后直接使用模型进行可视化结果测试:
torch.save(cnn_model.state_dict(), './cnnmodel')
